{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fa99bf4d",
   "metadata": {},
   "source": [
    "# 量化 LayerNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "571d2deb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.export import export\n",
    "import tvm\n",
    "from tvm import relax\n",
    "import tvm.testing\n",
    "from tvm.script import ir as I\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm.relax.frontend import detach_params\n",
    "from tvm.relax.frontend.torch import from_exported_program"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ffbabfa1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tuple(R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>dataflow():\n",
       "            lv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>layer_norm(x, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">0</span>], metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">1</span>], axes<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">2</span>, <span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">1</span>], epsilon<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1.0000000000000001e-05</span>, center<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>, scale<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "            gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tuple(R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">=</span> (lv,)\n",
       "            R<span style=\"color: #A2F; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "\n",
       "<span style=\"color: #007979; font-style: italic\"># Metadata omitted. Use show_meta=True in script() method to show it.</span>\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class LayerNorm(torch.nn.Module):\n",
    "    def __init__(self, shape):\n",
    "        super().__init__()\n",
    "        self.shape = shape\n",
    "        self.weight = torch.nn.Parameter(torch.ones(shape))\n",
    "        self.bias = torch.nn.Parameter(torch.zeros(shape))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.nn.functional.layer_norm(x, self.shape, self.weight, self.bias, 1e-5)\n",
    "\n",
    "example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)\n",
    "\n",
    "model = LayerNorm((10, 10))\n",
    "exported_program = export(model, args=example_args,)\n",
    "mod = from_exported_program(exported_program)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "432aff3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vm.builtin.check_tensor_info\n",
      "vm.builtin.check_tensor_info\n",
      "vm.builtin.match_shape\n",
      "vm.builtin.match_shape\n",
      "vm.builtin.alloc_storage\n",
      "vm.builtin.alloc_storage\n",
      "vm.builtin.alloc_tensor\n",
      "vm.builtin.alloc_tensor\n",
      "vm.builtin.null_value\n",
      "vm.builtin.null_value\n",
      "None 4\n",
      "layer_norm\n",
      "None 4\n",
      "layer_norm\n",
      "vm.builtin.make_tuple\n",
      "vm.builtin.make_tuple\n",
      "vm.builtin.null_value\n",
      "vm.builtin.null_value\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1947156/487285595.py:24: UserWarning: Returning type `vm.Storage` which is not registered via register_object, fallback to Object\n",
      "  output = vm[\"main\"](tvm.nd.array(data_np))\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import relax\n",
    "\n",
    "exe = tvm.compile(mod, \"llvm\")\n",
    "vm = relax.VirtualMachine(exe, tvm.cpu())\n",
    "hit_count = {}\n",
    "ret_vals = {}\n",
    "def instrument(func, name, before_run, ret_val, *args):\n",
    "    if (name, before_run) not in hit_count:\n",
    "        hit_count[(name, before_run)] = 0\n",
    "    hit_count[(name, before_run)] += 1\n",
    "    assert callable(func)\n",
    "    if before_run:\n",
    "        assert ret_val is None\n",
    "    if name == \"layer_norm\":\n",
    "        print(ret_val, len(args))\n",
    "    print(name)\n",
    "    if not before_run:\n",
    "        ret_vals[name] = ret_vals.get(name, []) + [ret_val]\n",
    "\n",
    "data_np = np.random.normal(size=(1, 3, 10, 10)).astype(\"float32\")\n",
    "vm.set_instrument(instrument)\n",
    "output = vm[\"main\"](tvm.nd.array(data_np))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd9a616b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['vm.builtin.check_tensor_info', 'vm.builtin.match_shape', 'vm.builtin.alloc_storage', 'vm.builtin.alloc_tensor', 'vm.builtin.null_value', 'layer_norm', 'vm.builtin.make_tuple'])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ret_vals.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "fd007bc7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1, 1, 1, 1, 2, 1, 1]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[len(v) for k, v in ret_vals.items()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "9461d612",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tvm.relax.vm_build.VMExecutable at 0x7ff42813ec60>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "exe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "2e3605fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ib = rx.Builder()\n",
      "with ib.function(\"main\", num_inputs=1):\n",
      "    ib.emit_call(\"vm.builtin.check_tensor_info\", args=[ib.r(0), ib.imm(4), ib.c(0), ib.c(1)])\n",
      "    ib.emit_call(\"vm.builtin.match_shape\", args=[ib.r(0), ib.r(18014398509481984), ib.imm(4), ib.imm(0), ib.imm(1), ib.imm(0), ib.imm(3), ib.imm(0), ib.imm(10), ib.imm(0), ib.imm(10), ib.c(1)])\n",
      "    ib.emit_call(\"vm.builtin.alloc_storage\", args=[ib.r(vm), ib.c(2), ib.imm(0), ib.c(3), ib.c(4)], dst=ib.r(1))\n",
      "    ib.emit_call(\"vm.builtin.alloc_tensor\", args=[ib.r(1), ib.imm(0), ib.c(5), ib.c(0)], dst=ib.r(2))\n",
      "    ib.emit_call(\"vm.builtin.null_value\", args=[], dst=ib.r(1))\n",
      "    ib.emit_call(\"layer_norm\", args=[ib.r(0), ib.c(6), ib.c(7), ib.r(2)])\n",
      "    ib.emit_call(\"vm.builtin.make_tuple\", args=[ib.r(2)], dst=ib.r(3))\n",
      "    ib.emit_call(\"vm.builtin.null_value\", args=[], dst=ib.r(2))\n",
      "    ib.emit_ret(ib.r(3))\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(exe.as_python())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "347fdcc8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
